# -*- coding: utf-8 -*-

import pandas as pd
from strategy import *
from database import eft_rt_data
from matplotlib import pyplot as plt
import numpy as np

pd.set_option("display.width", 300)

def macd_label(data_df, fast=9, slow=26, signal_period=12):
    """
    计算MACD买卖策略
    参数：
        df: 包含价格数据的DataFrame，需包含'close'列
        fast: 快线周期（默认12）
        slow: 慢线周期（默认26）
        signal_period: 信号线周期（默认9）
    返回：
        添加策略信号的DataFrame
    """
    df = data_df.copy()
    # 计算EMA
    df['EMA_fast'] = df['close'].ewm(span=fast, adjust=False).mean()
    df['EMA_slow'] = df['close'].ewm(span=slow, adjust=False).mean()

    # 计算MACD和信号线
    df['MACD'] = df['EMA_fast'] - df['EMA_slow']
    df['Signal'] = df['MACD'].ewm(span=signal_period, adjust=False).mean()

    # 生成交易信号（1: 买入， -1: 卖出， 0: 无操作）
    df['Signal_Cross'] = np.where(df['MACD'] > df['Signal'], 1, -1)
    df['Position'] = df['Signal_Cross'].diff()

    # 标记买卖点
    df['Buy_Signal'] = np.where((df['Position'] == -2), df['close'], 0)
    df['Sell_Signal'] = np.where((df['Position'] == 2), df['close'], 0)

    # 后处理
    bs_df = df[(df["Position"] == -2) | (df["Position"] == 2)].copy()
    print("==================macd====================")
    print(df[(df["Position"] == -2)].shape)
    print(df[(df["Position"] == 2)].shape)
    print("$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$")
    bs_df["price"] = bs_df["close"]
    bs_df["temp_signal"] = bs_df["Buy_Signal"].astype(bool).astype(int)
    bs_df.fillna(0, inplace=True)
    print(bs_df)
    return bs_df

def diff_label(data_df):
    df = data_df.copy()
    df = df[~(df["close"] == df["close"].shift(1))]
    # print(df["close"].ewm(span=5, adjust=True))
    df["diff"] = df["close"].diff(-1)
    df["diff1"] = (df["diff"] > 0).astype(int)  # 斜率为正
    df["diff2"] = - (df["diff"] < 0).astype(int)  # 斜率为负
    df["diff3"] = df["diff1"] + df["diff2"]
    df["Position"] = df["diff3"].diff(1)

    # 标记买卖点
    df['Buy_Signal'] = np.where((df['Position'] == -2), df['close'], 0)
    df['Sell_Signal'] = np.where((df['Position'] == 2), df['close'], 0)

    # 后处理
    bs_df = df[(df["Position"] == -2) | (df["Position"] == 2)].copy()
    print("=========================================macd====================")
    print(df[(df["Position"] == -2)].shape)
    print(df[(df["Position"] == 2)].shape)
    print("$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$")
    bs_df["price"] = bs_df["close"]
    bs_df["temp_signal"] = bs_df["Buy_Signal"].astype(bool).astype(int)
    bs_df.fillna(0, inplace=True)
    print(bs_df)
    return bs_df

def plot_graph(df):
    # print(df.head(50))
    fig = plt.figure()
    ax1 = fig.add_subplot(111)
    ax2 = ax1.twinx()
    ax1.plot(df["date"], df["close"])
    print(df.head(50))
    buy_df = df[~(df["price"].isnull()) & ~(df["temp_signal"] == 0)]
    sell_df = df[~(df["price"].isnull()) & (df["temp_signal"] == 0)]
    print("===============")
    print(buy_df)
    print(sell_df)
    ax2.scatter(buy_df["date"], buy_df["close"], s=10, c="r")
    ax2.scatter(sell_df["date"], sell_df["close"], s=15, c="g")
    ax1.set_xticklabels(df.date, rotation=90)
    plt.show()


if __name__ == "__main__":
    data_loader = eft_rt_data()
    df = data_loader.load_data(code=["513130"], start_date="20240801", end_date="20240803")
    # sign_df = diff_label(df)
    sign_df = macd_label(df)
    final_df = df.merge(sign_df[["price", "temp_signal"]], left_index=True, right_index=True, how="left")
    plot_graph(final_df)


